import torch.nn as nn
from .util import init, get_clones

"""MLP 模块。"""

class MLPLayer(nn.Module):
    """
    多层感知机层
    """
    def __init__(self, input_dim, hidden_size, layer_N, use_orthogonal, use_ReLU):
        """
        初始化多层感知机层
        
        参数:
            input_dim: 输入维度
            hidden_size: 隐藏层大小
            layer_N: 隐藏层数量
            use_orthogonal: 是否使用正交初始化
            use_ReLU: 是否使用ReLU激活函数
        """
        super(MLPLayer, self).__init__()
        self._layer_N = layer_N  # 隐藏层数量

        # 根据参数选择激活函数
        active_func = [nn.Tanh(), nn.ReLU()][use_ReLU]
        # 根据参数选择初始化方法
        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]
        # 计算增益
        gain = nn.init.calculate_gain(['tanh', 'relu'][use_ReLU])

        # 初始化函数
        def init_(m):
            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain=gain)

        # 第一个全连接层
        self.fc1 = nn.Sequential(
            init_(nn.Linear(input_dim, hidden_size)), active_func, nn.LayerNorm(hidden_size))
        # 隐藏层模板
        self.fc_h = nn.Sequential(init_(
            nn.Linear(hidden_size, hidden_size)), active_func, nn.LayerNorm(hidden_size))
        # 克隆多个隐藏层
        self.fc2 = get_clones(self.fc_h, self._layer_N)

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入张量
            
        返回:
            处理后的特征
        """
        # 通过第一个全连接层
        x = self.fc1(x)
        # 依次通过所有隐藏层
        for i in range(self._layer_N):
            x = self.fc2[i](x)
        return x


class MLPBase(nn.Module):
    """
    多层感知机基础网络，用作特征提取器
    """
    def __init__(self, args, obs_shape, cat_self=True, attn_internal=False):
        """
        初始化多层感知机基础网络
        
        参数:
            args: 参数对象
            obs_shape: 观察空间的形状
            cat_self: 是否连接自身特征
            attn_internal: 是否使用内部注意力
        """
        super(MLPBase, self).__init__()

        # 从参数中获取配置
        self._use_feature_normalization = args.use_feature_normalization  # 是否使用特征归一化
        self._use_orthogonal = args.use_orthogonal  # 是否使用正交初始化
        self._use_ReLU = args.use_ReLU  # 是否使用ReLU激活函数
        self._stacked_frames = args.stacked_frames  # 堆叠帧数
        self._layer_N = args.layer_N  # 隐藏层数量
        self.hidden_size = args.hidden_size  # 隐藏层大小

        # 获取观察空间维度
        obs_dim = obs_shape[0]

        # 如果使用特征归一化，创建层归一化
        if self._use_feature_normalization:
            self.feature_norm = nn.LayerNorm(obs_dim)

        # 创建多层感知机层
        self.mlp = MLPLayer(obs_dim, self.hidden_size,
                              self._layer_N, self._use_orthogonal, self._use_ReLU)

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入张量
            
        返回:
            处理后的特征
        """
        # 如果使用特征归一化，先进行归一化
        if self._use_feature_normalization:
            x = self.feature_norm(x)

        # 通过多层感知机层
        x = self.mlp(x)

        return x